package com.dong.es.entity.searchparam.agg;

import co.elastic.clients.elasticsearch._types.aggregations.*;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import com.dong.es.entity.EsBaseDTO;
import org.apache.commons.lang3.StringUtils;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @description: Es 聚合参数 去重计数
 * @author: 陈瑞东
 * @create: 2022-07-07 11:26
 */
public class EsAggsTermsParam extends EsAggsParam<List<EsAggsTermsParam.Result>> {
    private static final long serialVersionUID = -8275805738339507469L;

    // 字段名
    private String field;
    // 聚合大小
    private int size;
    // 类型
    private Class clazz;

    public EsAggsTermsParam(String field, Class clazz) {
        this(field, clazz, 65535);
    }

    public EsAggsTermsParam(String field, Class clazz, int size) {
        this.field = field;
        this.size = size;
        this.clazz = clazz;
        this.aggName = field + "Terms";
    }

    @Override
    protected boolean valid(EsBaseDTO esBaseDTO) {
        return StringUtils.isNotBlank(field) && size != 0;
    }

    @Override
    public void loadTo(SearchRequest.Builder target) {
        target.aggregations(this.aggName,
                       aggs -> aggs.terms(
                               terms -> terms
                                           .field(field)
                                           .size(size)
                       )
        );
    }

    @Override
    public List<Result> parseResult(Map<String, Aggregate> aggregations) {
        List<Result> result = new ArrayList<>();

        // 根据类型解析结果
        if (clazz == String.class) {
            List<StringTermsBucket> data = ((StringTermsAggregate) aggregations.get(this.aggName)._get()).buckets().array();
            data.forEach(e -> {
                result.add(new Result(e.key(), e.docCount()));
            });
        } else if (clazz == Long.class) {
            List<LongTermsBucket> data = ((LongTermsAggregate) aggregations.get(this.aggName)._get()).buckets().array();
            data.forEach(e -> {
                result.add(new Result(String.valueOf(e.key()), e.docCount()));
            });
        }

        return result;
    }

    /**
     * 聚合返回结果
     */
    public static class Result implements Serializable {
        private static final long serialVersionUID = -8952451162258868600L;
        // 名称
        private String name;
        // 数量
        private Long count;

        public Result(String name, Long count) {
            this.name = name;
            this.count = count;
        }

        public String getName() {
            return name;
        }

        public void setName(String name) {
            this.name = name;
        }

        public Long getCount() {
            return count;
        }

        public void setCount(Long count) {
            this.count = count;
        }

        @Override
        public String toString() {
            return "Result{" +
                    "name='" + name + '\'' +
                    ", count=" + count +
                    '}';
        }
    }
}
